import pandas as pd
import joblib
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score,classification_report
from xgboost import XGBClassifier
from utils.common import preprocess_data, plot_auc_curve


def train_model(data_path, model_dir, results_dir):
    """
    训练模型并评估[7](@ref)
    参数:
        data_path: 训练数据路径
        model_dir: 模型保存目录
        results_dir: 结果保存目录
    """
    # 1. 加载并预处理数据
    data = pd.read_csv(data_path)
    processed = preprocess_data(data)

    # 2. 分割特征和目标
    X = processed.drop(columns=['Attrition'])
    y = processed['Attrition']
    feature_columns = X.columns.tolist()

    # 3. 按8:2比例分割训练集和测试集[6](@ref)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    # 4. 数据标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 5. 训练模型
    model = XGBClassifier(
        n_estimators=100,
        max_depth=3,
        learning_rate=0.05,
        eval_metric='auc',
        random_state=42
    )
    model.fit(X_train_scaled, y_train)
    # print(" 模型训练完成")

    # 6. 评估训练集
    y_train_pred_proba = model.predict_proba(X_train_scaled)[:, 1]
    train_auc = roc_auc_score(y_train, y_train_pred_proba)
    # print(f" 训练集AUC分数: {train_auc:.4f}")
    print(classification_report(y_train, y_train_pred_proba > 0.5))



    # 7. 绘制AUC曲线
    plot_auc_curve(
        y_train, y_train_pred_proba,
        dataset_name="训练集",
        save_path=f'{results_dir}/train_auc_curve.png'
    )


    # 8. 保存模型和预处理对象
    joblib.dump(model, f'{model_dir}/xgb_model.pkl')
    joblib.dump(scaler, f'{model_dir}/scaler.pkl')
    joblib.dump(feature_columns, f'{model_dir}/feature_columns.pkl')
    # print(f" 模型已保存至 {model_dir}")

    return {"train_auc": train_auc}


if __name__ == "__main__":
    # 配置路径
    data_path = '../data/train.csv'
    model_dir = '../model'
    results_dir = '../data/fig'
    # 训练模型
    auc_scores = train_model(data_path, model_dir, results_dir)
    # print(" 训练完成:")
    print(f" 训练集AUC: {auc_scores['train_auc']:.4f}")
